import os
import numpy as np
import argparse
from ultralytics import YOLO
from PIL import Image
from paddleocr import PaddleOCR, draw_ocr
from PIL import Image, ImageDraw, ImageFont
from pdf2image import convert_from_path
import img2pdf
import fitz  
import cv2



def pdf_to_images(pdf_path, output_folder):
    """将PDF转换为图片"""
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    
    # 将PDF转换为图片
    images = convert_from_path(pdf_path)
    image_paths = []
    
    for i, image in enumerate(images):
        image_path = os.path.join(output_folder, f'page_{i:03d}.jpg')
        image.save(image_path, 'JPEG')
        image_paths.append(image_path)
    
    return image_paths

def process_ocr(image_path, font_path="C:/Windows/Fonts/simhei.ttf"):
    """处理OCR识别和可视化"""
    # 读取图片
    image = cv2.imread(image_path)
    
    # 使用PaddleOCR处理图片
    ocr = PaddleOCR(use_angle_cls=True, lang="ch")
    result = ocr.ocr(image, cls=True)
    
    # 准备OCR可视化
    image_pil = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
    draw = ImageDraw.Draw(image_pil)
    
    if result[0]:
        for line in result[0]:
            box = np.array(line[0], dtype=np.int32)
            text = line[1][0]
            
            # 计算文本框大小和调整字体
            text_height = max(point[1] for point in box) - min(point[1] for point in box)
            font_size = min(int(text_height * 0.8), 24)  # 增大最大字体大小
            font = ImageFont.truetype(font_path, font_size)
            
            # 绘制白色背景和文本
            draw.polygon([tuple(p) for p in box], fill=(255, 255, 255))
            x = min(point[0] for point in box)
            y = min(point[1] for point in box) + 2
            draw.text((x, y), text, font=font, fill=(0, 0, 0))
    
    return image_pil

def images_to_pdf(image_paths, output_pdf):
    """将图片合并为PDF"""
    with open(output_pdf, "wb") as f:
        f.write(img2pdf.convert(image_paths))

def main():
    # 设置路径
    pdf_path = r"D:\code\video-scan\resource\1.pdf"
    temp_folder = r"D:\code\video-scan\temp"
    output_folder = r"D:\code\video-scan\output"
    output_pdf = r"D:\code\video-scan\output\output_with_text.pdf"
    
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    
    try:
        # 1. PDF转换为图片
        print("正在将PDF转换为图片...")
        image_paths = pdf_to_images(pdf_path, temp_folder)
        
        # 2. 处理每张图片
        print("正在处理OCR识别...")
        processed_images = []
        for i, image_path in enumerate(image_paths):
            print(f"正在处理第 {i+1}/{len(image_paths)} 页")
            processed_image = process_ocr(image_path)
            output_path = os.path.join(output_folder, f'processed_page_{i:03d}.jpg')
            processed_image.save(output_path, 'JPEG', quality=95)
            processed_images.append(output_path)
        
        # 3. 合并为新的PDF
        print("正在生成最终PDF...")
        images_to_pdf(processed_images, output_pdf)
        
        print(f"处理完成！输出文件：{output_pdf}")
        
    except Exception as e:
        print(f"处理过程中出错: {str(e)}")

if __name__ == "__main__":
    main()